import os
import copy
import pickle
from numpy.lib.function_base import select
import pandas as pd
import datetime
import numpy as np
from datahelper import StockDataHelper

class FactorTest:
    def __init__(self,start_date: str, end_date: str):
        stockdatahelper = StockDataHelper()
        file_path = os.path.join(os.getcwd(),'..', 'data', 'data.pkl')
        stockdatahelper.init_data(file_path)
        # self.data_by_date = stockdatahelper.data_by_date
        self.data_by_date = stockdatahelper.get_start_end_dayData(start_date, end_date)
        self.trade_day_list = list(self.data_by_date.keys())

        self.data_by_stock = stockdatahelper.reform_data(self.data_by_date)
    def reform_data(self, dataset):
        """
        :param dataset: data.pkl的数据，字典{"trade_date":该trade_date所有股票的数据}
        :return: [['stock_id',该stock_id的股票数据dataframe],...]
        """
        if dataset is None:
            return None

        result = {}
        for date, info in dataset.items():
            info["trade_date"] = datetime.datetime.strptime(date, "%Y-%m-%d")
            for j in range(len(info.index)):
                stock, data = info.index[j], info.iloc[j: j + 1]
                if stock not in result:
                    result[stock] = data
                else:
                    result[stock] = result[stock].append(data)

        result = [[stock, info] for stock, info in result.items()]

        return result

    def get_period_returns(self, stock_df, start_date, end_date):
        """
        :start_date: str类型
        :end_date: str类型
        :return: start_date到end_date期间的该只股票的收益率，类型为一个浮点数，如-0.0036702428006776255
        """
        # str转为datetime，因为stock_df的trade_date列类型为datetime64
        # start_date = datetime.datetime.strptime(start_date, "%Y-%m-%d")  # 2010-01-04 00:00:00
        # end_date = datetime.datetime.strptime(end_date, "%Y-%m-%d")
        stock = copy.deepcopy(stock_df)
        stock = stock[(stock['trade_date'] >= start_date) & (stock['trade_date'] <= end_date)]

        stock = stock['close']  # 取close那一列
        try:
            period_return = stock.iloc[-1] / stock.iloc[0] - 1
            return period_return
        except:
            return np.nan

    def get_IC(self,factor_name):
        data_by_stock = self.data_by_stock
        trade_day_list = self.trade_day_list
        df_IC = pd.DataFrame(columns=['IC_pearson', 'IC_spearman'])
        for n in range(len(trade_day_list) - 1):
            trans_day = trade_day_list[n]
            next_trans_day = trade_day_list[n + 1]
            res = []
            # for i in range(len(data_by_stock)):
            for stock_id,stock_df in data_by_stock.items():
                # stock_id = data_by_stock[i][0]
                # stock_df = data_by_stock[i][1]

                period_return = self.get_period_returns(stock_df,trans_day, next_trans_day)

                factor = np.nan
                try:
                    stock_df = stock_df[stock_df['trade_date'] == trans_day]
                    factor = stock_df[factor_name].iloc[0]  # 取factor_name那一列
                except:
                    print("获取该股票的 %s Factor时出现问题" % factor_name)

                res_lst = [stock_id, period_return, factor]
                res.append(res_lst)

            df_factor = pd.DataFrame(res,
                                     columns=["stock", "period_return", factor_name]  # 指定列名
                                     )
            df_factor = df_factor.set_index("stock", inplace=False)  # 将stock列设为index

            # 计算因子的IC值
            df_factor = df_factor.dropna(how='any')  # 清理数据
            ic_pearson = df_factor.corr(method='pearson').iloc[1, 0]  # pearson IC
            ic_spearman = df_factor.corr(method='spearman').iloc[1, 0]  # spearman IC
            df_IC.loc[trans_day] = [ic_pearson, ic_spearman]

        ic_performance = pd.DataFrame(index=["IC_mean", "IC_std", "ICIR"], columns=["IC_pearson", "IC_spearman"])
        # IC 均值
        ic_performance.loc["IC_mean"] = df_IC.mean()
        # IC 标准差
        ic_performance.loc["IC_std"] = df_IC.std()
        # ICIR
        ic_performance.loc["ICIR"] = np.abs(df_IC.mean()) / df_IC.std()

        return df_IC, ic_performance

if __name__ == '__main__':
    start_date = '2010-01-04'
    end_date = '2010-01-08'
    factor_name = "MA"

    factortest = FactorTest(start_date,end_date)
    df_IC, ic_performance = factortest.get_IC(factor_name)
    print(df_IC)
    print(ic_performance)